"""
索引管理器

负责管理数据库中的所有索引，提供索引的创建、删除、查询等功能
"""

import os
import sys
import json
from typing import Dict, List, Optional, Any, Tuple

# 添加当前目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

try:
    from index_types import *
    from btree_index import BTreeIndex
    from statistics import Logger
    from constants import PAGE_SIZE
except ImportError:
    from .index_types import *
    from .btree_index import BTreeIndex
    from .statistics import Logger
    from .constants import PAGE_SIZE


class IndexManager:
    """索引管理器"""
    
    def __init__(self, page_manager, logger: Logger = None):
        self.page_manager = page_manager
        self.logger = logger or Logger()
        self.indexes = {}  # 索引名 -> BTreeIndex
        self.table_indexes = {}  # 表名 -> [索引名列表]
        self.index_metadata = {}  # 索引名 -> IndexMetadata
        self.metadata_page_id = None  # 索引元数据页面ID
        
        # 添加索引操作锁，确保线程安全
        import threading
        self._index_lock = threading.RLock()  # 使用递归锁，支持同一线程多次获取
        
        # 加载已存在的索引
        self._load_index_metadata()
    
    def create_primary_key_index(self, table_name: str, column_name: str, key_type: str = 'int') -> bool:
        """
        为表创建主键索引
        
        Args:
            table_name: 表名
            column_name: 主键列名
            key_type: 键类型 ('int', 'string', 'float')
        
        Returns:
            bool: 是否创建成功
        """
        index_name = f"pk_{table_name}_{column_name}"
        
        # 检查索引是否已存在
        if index_name in self.indexes:
            self.logger.warn("INDEX", 0, f"Primary key index {index_name} already exists")
            return False
        
        try:
            # 创建索引元数据
            metadata = IndexMetadata(
                index_name=index_name,
                table_name=table_name,
                column_name=column_name,
                index_type=IndexType.PRIMARY,
                key_type=key_type,
                root_page_id=0,  # 空树
                height=0,
                entry_count=0
            )
            
            # 创建B+树索引
            btree_index = BTreeIndex(metadata, self.page_manager, self.logger)
            
            # 注册索引
            self.indexes[index_name] = btree_index
            self.index_metadata[index_name] = metadata
            
            # 更新表索引映射
            if table_name not in self.table_indexes:
                self.table_indexes[table_name] = []
            self.table_indexes[table_name].append(index_name)
            
            # 保存元数据
            self._save_index_metadata()
            
            self.logger.info("INDEX", 0, f"Created primary key index {index_name} for table {table_name}")
            return True
            
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to create primary key index {index_name}: {e}")
            return False
    
    def create_secondary_index(self, table_name: str, column_name: str, index_name: str = None, 
                             unique: bool = False, key_type: str = 'int') -> bool:
        """
        创建二级索引
        
        Args:
            table_name: 表名
            column_name: 列名
            index_name: 索引名（可选，默认自动生成）
            unique: 是否为唯一索引
            key_type: 键类型
        
        Returns:
            bool: 是否创建成功
        """
        if index_name is None:
            index_name = f"idx_{table_name}_{column_name}"
        
        # 检查索引是否已存在
        if index_name in self.indexes:
            self.logger.warn("INDEX", 0, f"Index {index_name} already exists")
            return False
        
        try:
            # 创建索引元数据
            index_type = IndexType.UNIQUE if unique else IndexType.NORMAL
            metadata = IndexMetadata(
                index_name=index_name,
                table_name=table_name,
                column_name=column_name,
                index_type=index_type,
                key_type=key_type,
                root_page_id=0,
                height=0,
                entry_count=0
            )
            
            # 创建B+树索引
            btree_index = BTreeIndex(metadata, self.page_manager, self.logger)
            
            # 注册索引
            self.indexes[index_name] = btree_index
            self.index_metadata[index_name] = metadata
            
            # 更新表索引映射
            if table_name not in self.table_indexes:
                self.table_indexes[table_name] = []
            self.table_indexes[table_name].append(index_name)
            
            # 保存元数据
            self._save_index_metadata()
            
            self.logger.info("INDEX", 0, f"Created secondary index {index_name} for table {table_name}")
            return True
            
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to create secondary index {index_name}: {e}")
            return False
    
    def get_primary_key_index(self, table_name: str) -> Optional[BTreeIndex]:
        """获取表的主键索引"""
        if table_name not in self.table_indexes:
            return None
        
        for index_name in self.table_indexes[table_name]:
            metadata = self.index_metadata.get(index_name)
            if metadata and metadata.index_type == IndexType.PRIMARY:
                return self.indexes[index_name]
        
        return None
    
    def get_index(self, index_name: str) -> Optional[BTreeIndex]:
        """根据索引名获取索引"""
        return self.indexes.get(index_name)
    
    def get_table_indexes(self, table_name: str) -> List[BTreeIndex]:
        """获取表的所有索引"""
        indexes = []
        if table_name in self.table_indexes:
            for index_name in self.table_indexes[table_name]:
                if index_name in self.indexes:
                    indexes.append(self.indexes[index_name])
        return indexes
    
    def insert_to_indexes(self, table_name: str, record_data: Dict[str, Any], 
                         record_pointer: RecordPointer) -> bool:
        """向表的所有索引插入记录"""
        if table_name not in self.table_indexes:
            return True  # 没有索引，直接返回成功
        
        inserted_indexes = []  # 记录已插入的索引，用于回滚
        
        try:
            for index_name in self.table_indexes[table_name]:
                if index_name not in self.indexes or index_name not in self.index_metadata:
                    continue
                
                metadata = self.index_metadata[index_name]
                btree_index = self.indexes[index_name]
                
                # 获取该索引对应的列值
                column_value = record_data.get(metadata.column_name)
                if column_value is None:
                    # 处理NULL值（主键不能为NULL）
                    if metadata.index_type == IndexType.PRIMARY:
                        raise IndexException(f"Primary key column {metadata.column_name} cannot be NULL")
                    continue
                
                # 插入到索引
                btree_index.insert(column_value, record_pointer)
                inserted_indexes.append(index_name)
            
            return True
            
        except DuplicateKeyException as e:
            # 回滚已插入的索引
            self._rollback_index_inserts(inserted_indexes, record_data)
            self.logger.error("INDEX", 0, f"Duplicate key error during index insert: {e}")
            raise e
        except Exception as e:
            # 回滚已插入的索引
            self._rollback_index_inserts(inserted_indexes, record_data)
            self.logger.error("INDEX", 0, f"Failed to insert to indexes for table {table_name}: {e}")
            return False
    
    def delete_from_indexes(self, table_name: str, record_data: Dict[str, Any]) -> bool:
        """从表的所有索引删除记录"""
        with self._index_lock:
            if table_name not in self.table_indexes:
                return True
            
            success_count = 0
            total_count = 0
            
            for index_name in self.table_indexes[table_name]:
                if index_name not in self.indexes or index_name not in self.index_metadata:
                    continue
                
                metadata = self.index_metadata[index_name]
                btree_index = self.indexes[index_name]
                
                # 获取该索引对应的列值
                column_value = record_data.get(metadata.column_name)
                if column_value is None:
                    continue
                
                total_count += 1
                if btree_index.delete(column_value):
                    success_count += 1
            
            return success_count == total_count
    
    def search_by_primary_key(self, table_name: str, key_value: Any) -> Optional[RecordPointer]:
        """通过主键查找记录"""
        pk_index = self.get_primary_key_index(table_name)
        if pk_index is None:
            return None
        
        return pk_index.search(key_value)

    def delete_from_primary_key(self, table_name: str, key_value: Any) -> bool:
        """从主键索引删除指定键"""
        with self._index_lock:
            pk_index = self.get_primary_key_index(table_name)
            if pk_index is None:
                return False
            try:
                return pk_index.delete(key_value)
            except Exception as e:
                self.logger.warn("INDEX", 0, f"Primary key delete failed: {e}")
                return False

    def rebuild_primary_key_for_page(self, table_name: str, pk_column_name: str,
                                     page_id: int, records: list) -> None:
        """重建某一数据页上所有记录的主键→指针映射。

        注意：该方法会逐条 delete(key) 后再 insert(key, new_pointer)，
        以确保 slot_id 变更后索引中的指针保持一致。
        """
        pk_index = self.get_primary_key_index(table_name)
        if pk_index is None:
            return
        try:
            from index_types import RecordPointer
            for slot_id, rec in enumerate(records):
                key = None
                try:
                    key = rec.get_value(pk_column_name)
                except Exception:
                    continue
                if key is None:
                    continue
                # 先删除旧映射（若不存在会返回False，不影响）
                try:
                    pk_index.delete(key)
                except Exception:
                    pass
                # 插入新映射（指向新的 slot_id）
                try:
                    pk_index.insert(key, RecordPointer(page_id=page_id, slot_id=slot_id))
                except Exception as e:
                    # 若重复键异常，说明表中出现主键重复（不应发生），记录警告
                    self.logger.warn("INDEX", page_id, f"Rebuild PK mapping warning for key={key}: {e}")
        except Exception as e:
            self.logger.warn("INDEX", page_id, f"Rebuild PK page mappings failed: {e}")
    
    def search_by_index(self, index_name: str, key_value: Any) -> Optional[RecordPointer]:
        """通过索引名查找记录"""
        btree_index = self.get_index(index_name)
        if btree_index is None:
            return None
        
        return btree_index.search(key_value)
    
    def drop_table_indexes(self, table_name: str) -> bool:
        """删除表的所有索引"""
        if table_name not in self.table_indexes:
            return True
        
        try:
            index_names = self.table_indexes[table_name][:]
            for index_name in index_names:
                self.drop_index(index_name)
            
            del self.table_indexes[table_name]
            self._save_index_metadata()
            
            self.logger.info("INDEX", 0, f"Dropped all indexes for table {table_name}")
            return True
            
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to drop indexes for table {table_name}: {e}")
            return False
    
    def drop_index(self, index_name: str) -> bool:
        """删除指定索引"""
        if index_name not in self.indexes:
            return False
        
        try:
            metadata = self.index_metadata[index_name]
            table_name = metadata.table_name
            
            # 从内存中移除
            del self.indexes[index_name]
            del self.index_metadata[index_name]
            
            # 从表索引映射中移除
            if table_name in self.table_indexes:
                self.table_indexes[table_name].remove(index_name)
                if not self.table_indexes[table_name]:
                    del self.table_indexes[table_name]
            
            # 保存元数据
            self._save_index_metadata()
            
            # 释放索引占用的页面
            try:
                btree_index = self.indexes.get(index_name)
                if btree_index and hasattr(btree_index, 'release_pages'):
                    # 释放B+树索引占用的页面
                    btree_index.release_pages()
                    self.logger.info("INDEX", 0, f"Released pages for index {index_name}")
                elif btree_index and hasattr(btree_index, 'page_manager'):
                    # 如果有页面管理器，通过页面管理器释放
                    page_manager = btree_index.page_manager
                    if hasattr(page_manager, 'release_all_pages'):
                        page_manager.release_all_pages()
                        self.logger.info("INDEX", 0, f"Released pages via page manager for index {index_name}")
                else:
                    # 简化实现：标记索引相关页面为可回收
                    self.logger.info("INDEX", 0, f"Index {index_name} pages marked for garbage collection")
            except Exception as e:
                self.logger.warn("INDEX", 0, f"Failed to release pages for index {index_name}: {e}")
            
            self.logger.info("INDEX", 0, f"Dropped index {index_name}")
            return True
            
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to drop index {index_name}: {e}")
            return False
    
    def flush_all_indexes(self):
        """刷新所有索引的脏页"""
        with self._index_lock:
            for btree_index in self.indexes.values():
                try:
                    btree_index.flush_all()
                except Exception as e:
                    self.logger.error("INDEX", 0, f"Failed to flush index {btree_index.metadata.index_name}: {e}")
    
    def get_index_statistics(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        stats = {
            'total_indexes': len(self.indexes),
            'tables_with_indexes': len(self.table_indexes),
            'indexes': {}
        }
        
        for index_name, btree_index in self.indexes.items():
            stats['indexes'][index_name] = btree_index.get_statistics()
        
        return stats
    
    def _rollback_index_inserts(self, inserted_indexes: List[str], record_data: Dict[str, Any]):
        """回滚已插入的索引"""
        for index_name in inserted_indexes:
            try:
                metadata = self.index_metadata[index_name]
                btree_index = self.indexes[index_name]
                column_value = record_data.get(metadata.column_name)
                if column_value is not None:
                    btree_index.delete(column_value)
            except Exception as e:
                self.logger.error("INDEX", 0, f"Failed to rollback index insert for {index_name}: {e}")
    
    def _load_index_metadata(self):
        """加载索引元数据"""
        try:
            # 尝试查找索引元数据页面
            # 这里简化实现，假设元数据存储在固定的页面ID
            metadata_page_id = 1  # 假设页面1存储索引元数据
            
            try:
                metadata_data = self.page_manager.read_page(metadata_page_id)
                self.metadata_page_id = metadata_page_id
                
                # 解析元数据
                self._parse_metadata(metadata_data)
                
            except:
                # 元数据页面不存在，表示是新数据库
                self.metadata_page_id = None
                
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to load index metadata: {e}")
    
    def _parse_metadata(self, data: bytes):
        """解析索引元数据"""
        try:
            # 检查是否为空页面
            if data == b'\x00' * PAGE_SIZE:
                return
            
            # 解析JSON格式的元数据
            # 前4字节存储数据长度
            data_length = struct.unpack('<I', data[:4])[0]
            if data_length == 0 or data_length > PAGE_SIZE - 4:
                return
            
            json_data = data[4:4+data_length].decode('utf-8')
            metadata_dict = json.loads(json_data)
            
            # 重建索引
            for index_name, index_info in metadata_dict.items():
                metadata = IndexMetadata(
                    index_name=index_info['index_name'],
                    table_name=index_info['table_name'],
                    column_name=index_info['column_name'],
                    index_type=IndexType(index_info['index_type']),
                    key_type=index_info['key_type'],
                    root_page_id=index_info['root_page_id'],
                    height=index_info['height'],
                    entry_count=index_info['entry_count']
                )
                
                # 创建B+树索引对象
                btree_index = BTreeIndex(metadata, self.page_manager, self.logger)
                
                # 注册索引
                self.indexes[index_name] = btree_index
                self.index_metadata[index_name] = metadata
                
                # 更新表索引映射
                table_name = metadata.table_name
                if table_name not in self.table_indexes:
                    self.table_indexes[table_name] = []
                self.table_indexes[table_name].append(index_name)
            
            self.logger.info("INDEX", 0, f"Loaded {len(self.indexes)} indexes from metadata")
            
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to parse index metadata: {e}")
    
    def _save_index_metadata(self):
        """保存索引元数据"""
        try:
            # 构建元数据字典
            metadata_dict = {}
            for index_name, metadata in self.index_metadata.items():
                metadata_dict[index_name] = {
                    'index_name': metadata.index_name,
                    'table_name': metadata.table_name,
                    'column_name': metadata.column_name,
                    'index_type': metadata.index_type.value,
                    'key_type': metadata.key_type,
                    'root_page_id': metadata.root_page_id,
                    'height': metadata.height,
                    'entry_count': metadata.entry_count
                }
            
            # 序列化为JSON
            json_str = json.dumps(metadata_dict, indent=2)
            json_bytes = json_str.encode('utf-8')
            
            # 分配元数据页面（如果还没有）
            if self.metadata_page_id is None:
                self.metadata_page_id = self.page_manager.allocate_page()
            
            # 构建页面数据
            data_length = len(json_bytes)
            page_data = struct.pack('<I', data_length) + json_bytes
            
            # 填充到页面大小
            padding_size = PAGE_SIZE - len(page_data)
            if padding_size > 0:
                page_data += b'\x00' * padding_size
            else:
                # 数据太大，需要压缩或分页（简化实现暂不处理）
                self.logger.warn("INDEX", 0, "Index metadata too large for single page")
                page_data = page_data[:PAGE_SIZE]
            
            # 写入页面
            self.page_manager.write_page(self.metadata_page_id, page_data)
            
        except Exception as e:
            self.logger.error("INDEX", 0, f"Failed to save index metadata: {e}")
